import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.path import Path
import numpy as np

plt.rcParams.update({
    'font.sans-serif': 'Arial',
    'font.weight': 'bold',
    'mathtext.fontset': 'cm',
    'figure.facecolor': 'white',
    # 'text.usetex': True,  # 启用LaTeX
    # 'text.latex.preamble': r'\usepackage{amsmath}',  # 添加宏包
})

fig = plt.figure(figsize=(12, 6))
ax = fig.add_subplot(111)
ax.axis('off')

# 定义颜色方案
colors = {
    'cnn': '#4B8BBE',
    'rpn': '#FFD43B',
    'roi': '#72B352',
    'cls': '#C54B6C',
    'reg': '#9467BD',
    'arrow': '#307D7E'
}

# 左侧候选区域生成阶段
left_x = 0.05
# CNN特征提取
ax.add_patch(patches.Rectangle((left_x, 0.3), 0.15, 0.4, ec='k', fc=colors['cnn']))
plt.text(left_x+0.075, 0.5, 'CNN\nBackbone', ha='center', va='center', fontsize=10)

# RPN网络
rpn_x = left_x + 0.2
ax.add_patch(patches.Rectangle((rpn_x, 0.3), 0.15, 0.4, ec='k', fc=colors['rpn']))
plt.text(rpn_x+0.075, 0.5, 'RPN', ha='center', va='center', fontsize=12)

# 锚框矩阵
anchor_x = rpn_x + 0.2
for i in range(3):
    for j in range(3):
        ax.add_patch(patches.Rectangle(
            (anchor_x + i*0.07, 0.35 + j*0.1), 
            0.06, 0.08,
            ec='#FF0000', ls='--', fill=False, lw=1
        ))
plt.text(anchor_x+0.1, 0.68, 
         r'$w=128, h=64$'+'\n'+r'$\theta=0,\frac{\pi}{3}$',
         fontsize=8, color='red')

# 右侧分类回归阶段
right_x = 0.55
# ROI Pooling
ax.add_patch(patches.Rectangle((right_x, 0.3), 0.15, 0.4, ec='k', fc=colors['roi']))
plt.text(right_x+0.075, 0.5, 'ROI\nPooling', ha='center', va='center', fontsize=12)

# 分类器柱状图
cls_x = right_x + 0.2
categories = ['Car', 'Ped.', 'Cyc.', 'BG']
values = [0.7, 0.2, 0.08, 0.02]
for i, (cat, val) in enumerate(zip(categories, values)):
    ax.add_patch(patches.Rectangle(
        (cls_x, 0.3 + i*0.1), val*0.3, 0.08,
        fc=colors['cls'], ec='k'
    ))
    plt.text(cls_x + val*0.3 + 0.02, 0.34 + i*0.1, cat, fontsize=8)

# 边界框回归
reg_x = right_x + 0.35
ax.add_patch(patches.Rectangle((reg_x, 0.3), 0.15, 0.4, ec=colors['reg'], fc='none', lw=2))
plt.text(reg_x+0.075, 0.5, 
         r'$\Delta x=+0.2$'+'\n'+r'$\Delta y=-0.1$'+'\n'+
         r'$\Delta_w=+0.3$'+'\n'+r'$\Delta_h=+0.15$',
         fontsize=8, color=colors['reg'], ha='center')

# 损失函数标注
# 修改后的损失函数标注部分
plt.text(reg_x+0.08, 0.25, 
         r'$L_{\mathrm{loc}} = \sum_{i\in\{x,y,w,h\}} \mathrm{smooth}_{L1}(t_i)$',
         fontsize=9, rotation=30)

# 数据流箭头
arrow_style = dict(arrowstyle="->", color=colors['arrow'], lw=1.5)
plt.annotate('', xy=(left_x+0.15, 0.5), xytext=(left_x+0.2, 0.5),
             arrowprops=arrow_style)
plt.annotate('', xy=(rpn_x+0.15, 0.5), xytext=(anchor_x, 0.5),
             arrowprops=arrow_style)
plt.annotate('', xy=(anchor_x+0.3, 0.5), xytext=(right_x, 0.5),
             arrowprops=arrow_style)
plt.annotate('', xy=(right_x+0.15, 0.5), xytext=(cls_x-0.05, 0.5),
             arrowprops=arrow_style)

plt.tight_layout()
plt.savefig('two_stage_detector.svg', format='svg', bbox_inches='tight')
plt.show()
